import numpy as np
from database import VectorDatabase
from scipy.spatial.distance import cdist
# 初始化数据库
db = VectorDatabase()

# 模拟归一化后的特征向量
embeddings = np.random.rand(10, 512)
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)  # L2归一化

# 添加行人数据
person_ids = [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010]
db.add_person(embeddings, person_ids)

# 搜索相似行人（修复1：确保查询向量是二维的）
query = embeddings[3:4]  # 取第4个向量，保持形状 (1, 512)
distances = cdist(query, embeddings, metric="euclidean")
closest_idx = np.argmin(distances)
print(f"最小距离索引: {closest_idx}, 实际距离: {distances[0][closest_idx]}")
matches = db.search_person(query, top_k=1, threshold=2.0)  # 修复2：调整阈值
print(f"匹配结果: {matches}")  # 现在应该能匹配到 ['1004']

# 批量搜索（修复3：正确处理批量查询）
queries = embeddings[4:6]  # 形状 (2, 512)
batch_matches = db.search_person(queries, top_k=1, threshold=2.0)
print(f"批量匹配结果: {batch_matches}")  # 例如: [['1005'], ['1006']]

# 获取所有数据
all_data = db.get_all_persons()
print(f"总行人数量: {len(all_data['ids'])}")